import os
import logging
import argparse
import time
import json
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix

from app.ml.models import InvasionDetectionModel, load_and_preprocess_data

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 结果和模型保存目录
RESULTS_DIR = "results"
MODEL_DIR = "models"
os.makedirs(RESULTS_DIR, exist_ok=True)
os.makedirs(MODEL_DIR, exist_ok=True)


def train_model(model_type, dataset_path, test_size=0.2, random_state=42):
    """训练入侵检测模型
    
    Args:
        model_type: 模型类型，可选 "CNN", "LSTM", "CNN+LSTM"
        dataset_path: 数据集路径
        test_size: 测试集比例
        random_state: 随机种子
        
    Returns:
        dict: 训练结果和评估指标
    """
    start_time = time.time()
    
    logger.info(f"开始训练 {model_type} 模型")
    logger.info(f"数据集路径: {dataset_path}")
    
    # 1. 加载和预处理数据
    X, y, feature_names = load_and_preprocess_data(dataset_path)
    
    # 2. 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state, stratify=y
    )
    
    logger.info(f"训练集样本数: {X_train.shape[0]}, 测试集样本数: {X_test.shape[0]}")
    
    # 3. 创建模型
    model = InvasionDetectionModel(model_type=model_type)
    
    # 4. 训练模型 - 根据模型类型自动设置训练轮数
    history = model.train(X_train, y_train, epochs=None, batch_size=32, validation_split=0.2)
    
    # 5. 评估模型
    evaluation_results = model.evaluate(X_test, y_test)
    
    # 6. 保存模型
    model.save(MODEL_DIR)
    
    # 7. 构建结果字典
    training_time = time.time() - start_time
    
    results = {
        "model_type": model_type,
        "dataset": {
            "path": dataset_path,
            "n_samples": len(X),
            "n_features": X.shape[1],
            "n_classes": len(np.unique(y)),
            "train_test_split": {
                "test_size": test_size,
                "random_state": random_state
            }
        },
        "training": {
            "training_time": training_time
        },
        "evaluation": evaluation_results
    }
    
    # 8. 保存结果
    timestamp = time.strftime("%Y%m%d-%H%M%S")
    results_file = os.path.join(RESULTS_DIR, f"{model_type}_results_{timestamp}.json")
    with open(results_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    logger.info(f"训练完成，耗时 {training_time:.2f} 秒")
    logger.info(f"模型评估结果:")
    for metric, value in evaluation_results.items():
        if metric != "confusion_matrix":
            logger.info(f"  - {metric}: {value:.4f}")
    logger.info(f"结果已保存到 {results_file}")
    
    return results


def main():
    """命令行入口函数"""
    parser = argparse.ArgumentParser(description="训练入侵检测模型")
    parser.add_argument("--model", type=str, choices=["CNN", "LSTM", "CNN+LSTM"], 
                        default="CNN", help="模型类型")
    parser.add_argument("--dataset", type=str, required=True, help="数据集路径")
    parser.add_argument("--test_size", type=float, default=0.2, help="测试集比例")
    parser.add_argument("--random_state", type=int, default=42, help="随机种子")
    
    args = parser.parse_args()
    
    train_model(args.model, args.dataset, args.test_size, args.random_state)


if __name__ == "__main__":
    main() 